#include "lock.h"
#include "time.h"
#include "stdio.h"
#include "sched.h"
#include "queue.h"
#include "screen.h"
#include "irq.h"

pcb_t pcb[NUM_MAX_TASK];
/* current running task PCB */
pcb_t *current_running;

/* global process id */
pid_t process_id = 1;

/* kernel stack ^_^ */
#define NUM_KERNEL_STACK 20
#define STACK_BASE 0xffffffffa0f00000
#define STACK_OFFSET 0x10000

static uint64_t kernel_stack[NUM_KERNEL_STACK];
static int kernel_stack_count;

static uint64_t user_stack[NUM_KERNEL_STACK];
static int user_stack_count;

//int current_running_state = 0;

void init_stack()
{
     int i = 0;
     for (i = 0; i < NUM_KERNEL_STACK;i++)
     {
         kernel_stack[i] = STACK_BASE - 2 * i * STACK_OFFSET;
         user_stack[i]   = STACK_BASE - (2 * i + 1) * STACK_OFFSET;
     }
     kernel_stack_count = 0 ;
     user_stack_count = 0 ;
}

uint64_t new_kernel_stack()
{
     return kernel_stack[kernel_stack_count++];
}

uint64_t new_user_stack()
{
    return user_stack[user_stack_count++];
}

static void free_kernel_stack(uint64_t stack_addr)
{
     int i = 0;
     for (i = 0; i < kernel_stack_count; i++)
     {
         if(kernel_stack[i] == stack_addr)
            break;
     }
     for(; i<kernel_stack_count; i++)
     {
         kernel_stack[i] = kernel_stack[i+1] ;
     }
     kernel_stack[kernel_stack_count] = stack_addr;
     kernel_stack_count--;
}

static void free_user_stack(uint64_t stack_addr)
{
    int i = 0;
    for (i = 0; i < user_stack_count; i++)
    {
        if (user_stack[i] == stack_addr)
            break;
    }
    for (; i < user_stack_count; i++)
    {
        user_stack[i] = user_stack[i + 1];
    }
    user_stack[user_stack_count] = stack_addr;
    user_stack_count --;
}

/* Process Control Block */
void set_pcb(pid_t pid, pcb_t *pcb, task_info_t *task_info)
{
    memset(pcb->kernel_context.regs, 0, sizeof(uint64_t) * 32);
    pcb->which_queue = &ready_queue;
    pcb->lock_count = 0;
    pcb->pid = pid;
    pcb->type = task_info->type;
    pcb->status = TASK_READY;
    pcb->kernel_context.regs[31] = task_info->entry_point;
    pcb->kernel_context.cp0_epc = task_info->entry_point;
    pcb->kernel_context.cp0_status = 0x00008001;
    //pcb->kernel_context.cp0_cause = 0x40008000;
    pcb->user_context.regs[31] = task_info->entry_point;
    pcb->user_context.cp0_epc = task_info->entry_point;
    pcb->user_context.cp0_status = 0x00000001;
    //pcb->user_context.cp0_cause = 0x40008000;
    pcb->kernel_stack_top = new_kernel_stack();
    pcb->kernel_context.regs[29] = pcb->kernel_stack_top;
    pcb->user_stack_top = new_user_stack();
    pcb->user_context.regs[29] = pcb->user_stack_top;

    pcb->initial_priority = task_info->priority;
    pcb->wait_priority = pcb->initial_priority;
    pcb->counter = 0;

    pcb->kernel_state = 1;

    queue_init(&pcb->wait_queue);
}

/* ready queue to run */
queue_t ready_queue ;

/* block queue to wait */
queue_t block_queue ;

queue_t sleep_queue ;


static void check_sleeping()
{
    pcb_t *p = sleep_queue.head;
    pcb_t *q ;
    uint32_t time = get_timer();
    while (p != NULL)
    {
        if (time >= p->wake_time)
        {
            q = (pcb_t *)queue_remove(&sleep_queue, p);
            p->status = TASK_READY;
            ready_queue_push(&ready_queue, p);
            p = q;
        }
        else
            p = (void *)p->next;
    }
}

void scheduler(void)
{
    current_running->cursor_x = screen_cursor_x;
    current_running->cursor_y = screen_cursor_y;
    if(!queue_is_empty(&sleep_queue))
        check_sleeping();
    if (current_running->status != TASK_BLOCKED && current_running->status != TASK_EXITED)
    {
        current_running->status = TASK_READY;
        current_running->which_queue = &ready_queue;
        ready_queue_push(&ready_queue, current_running);
    }
    
    if(!queue_is_empty(&ready_queue))
    {
        current_running = (pcb_t *)queue_dequeue(&ready_queue);
        current_running->status = TASK_RUNNING;
        current_running->wait_priority = current_running->initial_priority;
        current_running->counter ++;
        current_running->which_queue = NULL;
    }
    if (current_running->status != TASK_BLOCKED && current_running->status != TASK_EXITED)
        current_running->status = TASK_RUNNING;
    if (!queue_is_empty(&ready_queue))
    {
        pcb_t *p = (pcb_t *)ready_queue.head;
        while(p != NULL)
        {
            p->wait_priority ++;
            p = (pcb_t*) p-> next;
        }
    }
    //vt100_move_cursor(1, 27+current_running->pid);
    //printk("> [TASK] This taaaaaask (pid = %d) has been scheduled %d times", current_running->pid,current_running->counter);
    if (current_running->type == USER_PROCESS || current_running->type == USER_THREAD)
        current_running->kernel_state = 0;
    else
        current_running->kernel_state = 1;
    current_running->which_queue = NULL;
    screen_cursor_x = current_running->cursor_x;
    screen_cursor_y = current_running->cursor_y;
}

void do_sleep(uint32_t sleep_time)
{
     current_running -> status = TASK_BLOCKED;
     current_running -> wake_time = get_timer() + sleep_time;
     current_running -> which_queue = &sleep_queue;
     queue_push(&sleep_queue,(void*)current_running);

     do_scheduler();
}

void do_exit(void)
{
    pcb_t *pcb_exit = current_running;
    int i;
    for (i = 0; pcb[i].pid != current_running->pid; i++);
    pcb[i].pid = -1;
    pcb_exit->status = TASK_EXITED;
    pcb_exit->which_queue = NULL;
    do_unblock_all(&pcb_exit->wait_queue);
    for (i = 0; i < pcb_exit->lock_count; i++)
        do_mutex_lock_release(pcb_exit->lock[i]);
    free_user_stack(pcb_exit->user_stack_top);
    free_kernel_stack(pcb_exit->kernel_stack_top);
    do_scheduler();
}

void do_block(queue_t *queue)
{
    current_running->status = TASK_BLOCKED;
    current_running->which_queue = queue;
    queue_push(queue,current_running);
}

void do_unblock_one(queue_t *queue)
{
     pcb_t *unblock_one;
     unblock_one = queue_dequeue(queue);
     unblock_one->status = TASK_READY;
     unblock_one->which_queue = &ready_queue;
     ready_queue_push(&ready_queue,unblock_one);
}

void do_unblock_all(queue_t *queue)
{
    pcb_t *unblock_one;
    while(!queue_is_empty(queue))
    {
        unblock_one = queue_dequeue(queue);
        unblock_one->status = TASK_READY;
        unblock_one->which_queue = &ready_queue;
        ready_queue_push(&ready_queue, unblock_one);
    }
}

int do_spawn(task_info_t *task)
{
    pcb_t *pcb_exec;
    int i = 0;
    for(i = 0; pcb[i].pid != -1; i++);
    pcb_exec = &pcb[i];
    set_pcb(process_id++, pcb_exec, task);
    ready_queue_push(&ready_queue, (void*)pcb_exec);
}

int do_kill(pid_t pid)
{
    int i;
    for (i = 0; pcb[i].pid != pid; i++);
    pcb_t *pcb_kill = &pcb[i];

    if (pcb_kill->which_queue != NULL && pcb_kill != current_running)
        queue_remove(pcb_kill->which_queue, (void *)pcb_kill);
    pcb[i].pid = -1;
    pcb_kill->status = TASK_EXITED;
    pcb_kill->which_queue = NULL;
    do_unblock_all(&pcb_kill->wait_queue);

    for (i = 0; i < pcb_kill->lock_count; i++)
        do_mutex_lock_release(pcb_kill->lock[i]);
    free_user_stack(pcb_kill->user_stack_top);
    free_kernel_stack(pcb_kill->kernel_stack_top);

    if(pcb_kill == current_running)
        do_scheduler();
}

int do_waitpid(pid_t pid)
{
    int i;
    for (i = 0; pcb[i].pid != pid; i++);
    pcb_t *pcb_wait = &pcb[i];
    if (pcb_wait->status != TASK_EXITED)
    {
        current_running->status = TASK_BLOCKED;
        current_running->which_queue = &pcb_wait->wait_queue;
        queue_push(&pcb_wait->wait_queue, (void *)current_running);
        do_scheduler();
    }
}

// process show
void do_process_show()
{
    kprintf("[PROCESS TABLE]\n");
    int i = 0, j = 0;
    for (i = 0; i < NUM_MAX_TASK; i++)
    {
        if (pcb[j].status == TASK_RUNNING)
        {
            kprintf("[%d] PID: %d STATUS: RUNNING\n", j, pcb[i].pid);
        }
        else if (pcb[i].status == TASK_BLOCKED )
        {
            kprintf("[%d] PID: %d STATUS: BLOCKED\n", j, pcb[i].pid);
        }
        else if (pcb[j].status == TASK_READY)
        {
            kprintf("[%d] PID: %d STATUS: READY\n", j, pcb[j].pid);
        }
        else
        {
            continue;
        }
        j ++;
        
    }
}

pid_t do_getpid()
{
    return current_running->pid;
}